package BSTChallengesPack;

import TreeNodePack.TreeNode;

import java.util.ArrayList;
import java.util.List;

public class BSTChallenges {

    // 1. 找到中序后继节点
    public TreeNode inorderSuccessor(TreeNode root, TreeNode p) {
        // 情况1：p有右子树，后继是右子树的最小节点
        if (p.right != null) {
            return findMin(p.right);
        }
        // 情况2：从根开始寻找最近的比p大的祖先
        TreeNode successor = null;
        TreeNode current = root;
        while (current != null) {
            if (current.val > p.val) {
                successor = current;
                current = current.left;
            } else {
                current = current.right;
            }
        }
        return successor;
    }

    private TreeNode findMin(TreeNode node) {
        while (node.left != null) {
            node = node.left;
        }
        return node;
    }

    // 2. 计算区间和
    public int rangeSumBST(TreeNode root, int low, int high) {
        if (root == null) return 0;
        // 当前节点值大于high，只需处理左子树
        if (root.val > high) {
            return rangeSumBST(root.left, low, high);
        }
        // 当前节点值小于low，只需处理右子树
        if (root.val < low) {
            return rangeSumBST(root.right, low, high);
        }
        // 当前节点在区间内，累加并递归左右子树
        return root.val + rangeSumBST(root.left, low, high) + rangeSumBST(root.right, low, high);
    }

    // 3. 合并两棵BST的所有元素
    public List<Integer> getAllElements(TreeNode root1, TreeNode root2) {
        List<Integer> list1 = new ArrayList<>();
        List<Integer> list2 = new ArrayList<>();
        inorder(root1, list1);
        inorder(root2, list2);
        return mergeLists(list1, list2);
    }

    private void inorder(TreeNode root, List<Integer> list) {
        if (root == null) return;
        inorder(root.left, list);
        list.add(root.val);
        inorder(root.right, list);
    }

    private List<Integer> mergeLists(List<Integer> list1, List<Integer> list2) {
        List<Integer> result = new ArrayList<>();
        int i = 0, j = 0;
        while (i < list1.size() && j < list2.size()) {
            if (list1.get(i) < list2.get(j)) {
                result.add(list1.get(i++));
            } else {
                result.add(list2.get(j++));
            }
        }
        while (i < list1.size()) result.add(list1.get(i++));
        while (j < list2.size()) result.add(list2.get(j++));
        return result;
    }

    // 测试代码
    public static void main(String[] args) {
        BSTChallenges challenges = new BSTChallenges();

        // 测试中序后继
        TreeNode root = new TreeNode(5);
        root.left = new TreeNode(3);
        root.right = new TreeNode(7);
        root.left.left = new TreeNode(2);
        root.left.right = new TreeNode(4);
        root.right.left = new TreeNode(6);
        root.right.right = new TreeNode(8);

        TreeNode p = root.left.right; // 节点4
        TreeNode successor = challenges.inorderSuccessor(root, p);
        System.out.println("Inorder successor of 4: " + (successor != null ? successor.val : "null")); // 5

        // 测试区间和
        int sum = challenges.rangeSumBST(root, 3, 7);
        System.out.println("Range sum [3,7]: " + sum); // 3+4+5+6+7=25

        // 测试合并两棵BST
        TreeNode root2 = new TreeNode(1);
        root2.right = new TreeNode(9);
        List<Integer> mergedList = challenges.getAllElements(root, root2);
        System.out.println("Merged list: " + mergedList); // [1,2,3,4,5,6,7,8,9]
    }
}